# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
#

from scipy import zeros, ones, linspace, asarray
from scipy import arange, linalg

class Grid:
    """A simple grid class that stores the details and solution of the
    computational grid."""

    __solts__=['m_nx', 'm_ny', 'm_xmin', 'm_xmax', 'm_ymin', 'm_ymax', 'm_u', 'm_old_u']

    def __init__(self, nx=10, ny=10, xmin=0.0, xmax=1.0,
                 ymin=0.0, ymax=1.0):
        self.m_xmin, self.m_xmax, self.m_ymin, self.m_ymax = xmin, xmax, ymin, ymax
        self.m_nx = nx
        self.m_ny = ny
        self.m_dx = float(xmax-xmin)/(nx-1)
        self.m_dy = float(ymax-ymin)/(ny-1)
        self.m_u  = zeros((nx, ny), 'd')

        # used to compute the change in solution in some of the methods.
        self.m_old_u = self.m_u.copy()

    def setBCFunc(self, func):
      """
        When use as a finite difference mesh, Sets the BC given a function of two variables.
      """
      m_xmin, ymin = self.m_xmin, self.m_ymin
      xmax, ymax = self.m_xmax, self.m_ymax

      x = arange(m_xmin, xmax + self.m_dx*0.5, self.m_dx)
      y = arange(ymin, ymax + self.m_dy*0.5, self.m_dy)

      self.m_u[0 ,:] = func(m_xmin, y)
      self.m_u[-1,:] = func(xmax, y)
      self.m_u[:, 0] = func(x, ymin)
      self.m_u[:,-1] = func(x, ymax)

    def get_grid(self):
      return self.m_u

    def getBox(self):
      return [self.m_xmin, self.m_xmax, self.m_ymin, self.m_ymax]

    def getDeltaXY(self):
      return [self.m_dx, self.m_dy]

    def getNxNy(self):
      return [self.m_nx, self.m_ny]

    def lexPositions(self, i, j):
      """ return the lexicographic position of point (i, j)
          Dirichlet BC.
      """
      assert( 0 < i and i < self.nx )
      assert( 0 < j and j < self.ny )
      return (self.nx-2)*(i-1)+(j-1)

    def computeError(self):
     """Computes absolute error using an L2 norm for the solution.
        This requires that self.m_u and self.m_m_old_u must be appropriately
        setup.
     """
     v = (self.m_u - self.m_old_u).flat
     return linalg.norm(v)

    def getTriangle(self, i, j, posi):
      """
        returns the coordinates of the triangle
        with appex (i, j) and in the posi quadrant.
        posi \in [0, 3]
        0 ~ up | up
        1 ~ up | do
        2 ~ do | do
        3 ~ do | do
      """

      assert( 0 <= i and i < self.nx )
      assert( 0 <= j and j < self.ny )
      assert( 0 <= posi and posi < 4 )

      neig ={ 0 : [1, 1],
              1 : [-1, 1],
              2 : [-1, -1],
              3 : [1, -1]
            }[posi]

      if 0 > i+neig[0] or i+neig[0] > self.m_nx:
        return None

      if 0 > j+neig[1] or j+neig[1] > self.m_ny:
        return None

      dx = self.m_dx
      dy = self.m_dy

      tria = zeros( (3, 2), 'float')
      ll   = asarray([self.m_xmin, self.ymin])
      tria[0, :] = asarray( [i*dx, j*dy] + ll )

      if 0 == posi:
        tria[1, :] = tria[0, :] + asarray([dx, 0])
        tria[2, :] = tria[0, :] + asarray([0 , dy])
      if 1 == posi:
        tria[1, :] = tria[0, :] + asarray([0  , dy])
        tria[2, :] = tria[0, :] + asarray([-dx, 0])
      if 2 == posi:
        tria[1, :] = tria[0, :] + asarray([-dx, 0])
        tria[2, :] = tria[0, :] + asarray([0, -dy])
      if 3 == posi:
        tria[1, :] = tria[0, :] + asarray([0  , -dy])
        tria[2, :] = tria[0, :] + asarray([-dx, 0])

      return tria

class Grid1D(object):
    __solts__=['m_nSegs', 'm_x', 'm_xmin', 'm_xmax']
    def __init__(self, nSegs = 10, xmin=0.0, xmax=1.0):

        self.m_nSegs = nSegs
        self.m_x  = asarray(linspace(xmin, xmax, nSegs+1, 'True'))
        self.m_dx = float(xmax-xmin)/nSegs
        (self.m_xmin, self.m_xmax) = (xmin, xmax)

    def get_nsegs(self):
        return self.m_nSegs

    def getSegments(self):
        return [asarray([self.m_x[i], self.m_x[i+1]], 'float') for i in range(self.m_nSegs)]

    def getDx(self):
        return self.m_x[1:]-self.m_x[:-1]

    def getNumSegs(self):
        return self.m_nSegs

    def get_grid(self):
        return self.m_x

    def getRange(self):
        return (self.m_xmin, self.m_xmax)




